from peewee import *
import datetime

# 创建数据库连接（使用SQLite作为示例，可以轻松切换到MySQL等）
database = SqliteDatabase('chat_history.db')

class BaseModel(Model):
    """基础模型类"""
    class Meta:
        database = database

class Conversation(BaseModel):
    """对话记录模型"""
    session_id = CharField(max_length=100, index=True)  # 会话ID
    role = CharField(max_length=20)  # 角色：user 或 assistant
    content = TextField()  # 消息内容
    timestamp = DateTimeField(default=datetime.datetime.now)  # 时间戳
    created_at = DateTimeField(default=datetime.datetime.now)  # 创建时间
    
    class Meta:
        table_name = 'conversations'
        indexes = (
            (('session_id', 'timestamp'), False),  # 复合索引
        )

class ChatSession(BaseModel):
    """聊天会话模型"""
    session_id = CharField(max_length=100, unique=True, index=True)  # 会话ID
    user_id = CharField(max_length=100, null=True)  # 用户ID（可选）
    title = CharField(max_length=200, null=True)  # 会话标题
    created_at = DateTimeField(default=datetime.datetime.now)  # 创建时间
    updated_at = DateTimeField(default=datetime.datetime.now)  # 更新时间
    message_count = IntegerField(default=0)  # 消息数量
    
    class Meta:
        table_name = 'chat_sessions'

def create_tables():
    """创建数据库表"""
    database.connect()
    database.create_tables([Conversation, ChatSession], safe=True)
    database.close()

def get_database():
    """获取数据库连接"""
    return database 